#! /usr/bin/env python3
import cv2
import numpy as np
import os
from glob import glob
from PIL import Image
from skimage.transform import hough_line, hough_line_peaks, rotate
import argparse
try:
    from Preprocessing.libs.datasetFuncs import formatImage
except ImportError as e:
    from libs.datasetFuncs import formatImage

from keras.models import load_model
current_path = os.path.split(os.path.abspath(__file__))[0]
model = load_model(os.path.join(current_path, './textdeskew.h5'))

format_size = (128, 128)
# K = np.sum(np.array(format_size)**2)**0.5

def calculate_deviation(angle):
    angle_in_degrees = np.abs(angle)
    deviation = np.abs(np.pi / 4 - angle_in_degrees)

    return deviation

def compare_sum(value):
    if value >= 44 and value <= 46:
        return True
    else:
        return False

def get_max_freq_elem(arr):
    max_arr = []
    freqs = {}
    for i in arr:
        if i in freqs:
            freqs[i] += 1
        else:
            freqs[i] = 1

    sorted_keys = sorted(freqs, key=freqs.get, reverse=True)
    max_freq = freqs[sorted_keys[0]]

    for k in sorted_keys:
        if freqs[k] == max_freq:
            max_arr.append(k)

    return max_arr

def textDeskew(input_image, output_path, limit=1, no_bias=False):
    if isinstance(input_image, Image.Image):
        img = np.array(input_image, dtype='uint8')
    else:
        img = cv2.imread(input_image)

    if len(img.shape) == 3:
        gray = img.mean(axis=2, dtype='uint8') # cv2.imread(input_image, 0)
    else:
        gray = img

    bias_angle = 0
    if not no_bias:
        size = np.array(gray.shape[::-1])
        sample_k = 2048 / size.max()
        sample_size = np.around(sample_k * size).astype('int')
        sample = cv2.resize(gray, tuple(sample_size), interpolation=cv2.INTER_AREA)

        edges = cv2.Canny(sample, 50, 150, apertureSize=3)

        h, a, d = hough_line(edges)
        _, ap, _ = hough_line_peaks(h, a, d, num_peaks=20)

        if len(ap) != 0:
            absolute_deviations = [calculate_deviation(k) for k in ap]
            average_deviation = np.mean(np.rad2deg(absolute_deviations))
            ap_deg = [np.rad2deg(x) for x in ap]

            bin_0_45 = []
            bin_45_90 = []
            bin_0_45n = []
            bin_45_90n = []
            
            for ang in ap_deg:
                deviation_sum = int(90 - ang + average_deviation)
                if compare_sum(deviation_sum):
                    bin_45_90.append(ang)
                    continue

                deviation_sum = int(ang + average_deviation)
                if compare_sum(deviation_sum):
                    bin_0_45.append(ang)
                    continue

                deviation_sum = int(-ang + average_deviation)
                if compare_sum(deviation_sum):
                    bin_0_45n.append(ang)
                    continue

                deviation_sum = int(90 + ang + average_deviation)
                if compare_sum(deviation_sum):
                    bin_45_90n.append(ang)

            angles = [bin_0_45, bin_45_90, bin_0_45n, bin_45_90n]
            lmax = 0

            for j in range(len(angles)):
                l = len(angles[j])
                if l > lmax:
                    lmax = l
                    maxi = j

            if lmax:
                ans_arr = get_max_freq_elem(angles[maxi])
                ans_res = np.mean(ans_arr)

            else:
                ans_arr = get_max_freq_elem(ap_deg)
                ans_res = np.mean(ans_arr)

            angle = ans_res

            if angle >= 0 and angle <= 90:
                bias_angle = angle - 90
            if angle >= -45 and angle < 0:
                bias_angle = angle - 90
            if angle >= -90 and angle < -45:
                bias_angle = 90 + angle

        rotated1 = rotate(img, bias_angle, resize=True, cval=1)
    else:
        rotated1 = img.astype('float') / 255

    mini = np.ones(format_size)
    rotated1_size = np.array([rotated1.shape[1], rotated1.shape[0]])
    pool_k = 128 / rotated1_size.max()
    pool_size = np.around(pool_k * rotated1_size).astype('int')
    pool = cv2.resize(rotated1, tuple(pool_size), interpolation=cv2.INTER_AREA)
    if len(pool.shape) > 2: pool = pool.mean(axis=2)
    mini_limits = (pool_size - pool_size.min()) // 2
    mini[mini_limits[0]: mini_limits[0]+pool_size[1], mini_limits[1]: mini_limits[1]+pool_size[0]] = pool
    input_X = 1 - mini.reshape(-1, format_size[0], format_size[1], 1).astype('float32')
    predict_probs = model.predict(input_X)
    direction_num = (predict_probs == predict_probs.max()).dot(np.arange(4))

    rotated2 = (np.rot90(rotated1, k=-direction_num) * 255).astype('uint8')

    cv2.imwrite(output_path, rotated2)

    print('{}: Deskew Angle: {:.2f}'.format(input_image, bias_angle + -90 * int(direction_num)))
    return rotated2

def textDeskews(input_dir, output_dir, no_bias=False):
    input_paths = glob(os.path.join(input_dir, '*'))
    for input_path in input_paths:
        basename = os.path.basename(input_path)
        output_path = os.path.join(output_dir, basename)
        textDeskew(input_path, output_path, no_bias=no_bias)

if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='深度学习图片矫偏工具')
    parser.add_argument('input_dir', metavar='原始图片目录', type=str, help='原始图片所在目录')
    parser.add_argument('output_dir', metavar='生成图片目录', type=str, help='生成图片所在目录')
    parser.add_argument('--no_bias', dest='no_bias', nargs='?', const=True, default=False,
                        help='指定是否取消偏置调整，若不启用则采用偏置调整')

    args = parser.parse_args()
    textDeskews(args.input_dir, args.output_dir, no_bias=args.no_bias)
